"""
简化版A股股票预测系统 - 重点优化股票代码输入
"""

import streamlit as st
import sys
import os
import pandas as pd

# 添加src目录到路径
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))

from data_loader import get_stock_data, add_technical_indicators, get_stock_info, validate_stock_code
from predictor import StockPredictor
from visualizer import plot_stock_price, plot_technical_indicators, plot_prediction_results, display_metrics

# 页面配置
st.set_page_config(
    page_title="A股股票预测系统",
    page_icon="📈",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 主标题
st.title("📈 A股股票预测系统")
st.markdown("**通过股票代码快速获取任意A股数据进行预测分析**")

# 核心功能：股票代码输入
st.markdown("---")
col1, col2, col3 = st.columns([2, 1, 1])

with col1:
    st.markdown("### 🎯 输入股票代码")
    
    # 主要输入框
    stock_input = st.text_input(
        "请输入6位股票代码",
        placeholder="例如: 600519(茅台), 000001(平安), 300750(宁德时代)",
        help="支持所有A股代码，输入6位数字即可"
    )
    
    # 实时验证
    if stock_input:
        is_valid, stock_info = validate_stock_code(stock_input)
        if is_valid:
            st.success(f"✅ 找到股票: {stock_info.get('name', '未知股票')} ({stock_info['code']})")
            current_stock_code = stock_info['code']
            current_stock_name = stock_info.get('name', f"股票{stock_info['code']}")
        else:
            st.error("❌ 请输入正确的6位股票代码")
            current_stock_code = None
            current_stock_name = None
    else:
        current_stock_code = None
        current_stock_name = None

with col2:
    st.markdown("### 🔥 热门推荐")
    hot_stocks = [
        ("茅台", "600519"),
        ("平安", "601318"), 
        ("招行", "600036"),
        ("比亚迪", "002594"),
        ("宁德", "300750")
    ]
    
    for name, code in hot_stocks:
        if st.button(f"{name}\n{code}", key=f"hot_{code}", use_container_width=True):
            st.session_state.selected_stock = {'name': name, 'code': code}
            st.rerun()

with col3:
    st.markdown("### ⚙️ 预测设置")
    predict_days = st.slider("预测天数", 5, 30, 20)
    data_days = st.slider("历史数据(天)", 365, 1095, 730)

# 检查是否有选择的股票
if 'selected_stock' in st.session_state:
    current_stock_code = st.session_state.selected_stock['code']
    current_stock_name = st.session_state.selected_stock['name']
    st.info(f"🎯 当前选择: {current_stock_name} ({current_stock_code})")

# 数据加载和分析
if current_stock_code and current_stock_name:
    
    # 数据加载按钮
    if st.button(f"🔄 加载 {current_stock_name} 数据并开始分析", type="primary", use_container_width=True):
        
        with st.spinner(f"正在获取 {current_stock_name} 数据..."):
            try:
                # 获取数据
                df, data_type = get_stock_data(current_stock_code, current_stock_name, data_days)
                df = add_technical_indicators(df)
                
                st.success(f"✅ 成功获取 {len(df)} 条{data_type}")
                
                # 存储到session state
                st.session_state.df = df
                st.session_state.current_stock = {'name': current_stock_name, 'code': current_stock_code}
                
            except Exception as e:
                st.error(f"❌ 数据获取失败: {str(e)}")

# 显示分析结果
if 'df' in st.session_state and st.session_state.df is not None:
    df = st.session_state.df
    stock_info = get_stock_info(df, st.session_state.current_stock['name'])
    
    st.markdown("---")
    st.markdown(f"## 📊 {st.session_state.current_stock['name']} 分析结果")
    
    # 基本信息
    col1, col2, col3, col4 = st.columns(4)
    
    with col1:
        st.metric("💰 当前价格", f"¥{stock_info['current_price']:.2f}",
                 f"{stock_info['change']:+.2f} ({stock_info['change_pct']:+.2f}%)")
    
    with col2:
        st.metric("📈 52周最高", f"¥{stock_info['high_52w']:.2f}")
    
    with col3:
        st.metric("📉 52周最低", f"¥{stock_info['low_52w']:.2f}")
    
    with col4:
        if 'rsi' in stock_info:
            rsi_status = "超买" if stock_info['rsi'] > 70 else "超卖" if stock_info['rsi'] < 30 else "正常"
            st.metric("📊 RSI", f"{stock_info['rsi']:.1f}", rsi_status)
    
    # 图表展示
    tab1, tab2, tab3 = st.tabs(["📈 价格走势", "📊 技术指标", "🔮 智能预测"])
    
    with tab1:
        price_chart = plot_stock_price(df, f"{st.session_state.current_stock['name']}股票走势")
        st.plotly_chart(price_chart, use_container_width=True)
    
    with tab2:
        tech_chart = plot_technical_indicators(df)
        st.plotly_chart(tech_chart, use_container_width=True)
        
        # 最新数据
        st.subheader("📋 最新数据")
        st.dataframe(
            df.tail(5)[['open', 'high', 'low', 'close', 'volume', 'ma5', 'ma20', 'rsi']].round(2),
            use_container_width=True
        )
    
    with tab3:
        col1, col2 = st.columns(2)
        
        with col1:
            if st.button("🧠 训练预测模型", type="primary", use_container_width=True):
                with st.spinner("训练中..."):
                    try:
                        predictor = StockPredictor(sequence_length=10)
                        train_results = predictor.train(df, test_ratio=0.2)
                        st.session_state.predictor = predictor
                        st.session_state.model_trained = True
                        st.success("✅ 模型训练完成！")
                        display_metrics(train_results['test_metrics'], "🎯 模型性能")
                    except Exception as e:
                        st.error(f"❌ 训练失败: {str(e)}")
        
        with col2:
            if st.session_state.get('model_trained', False):
                if st.button(f"🚀 预测未来{predict_days}天", type="secondary", use_container_width=True):
                    with st.spinner("预测中..."):
                        try:
                            predictor = st.session_state.predictor
                            results = predictor.predict_future(df, days=predict_days)
                            st.session_state.prediction_results = results
                            st.success("✅ 预测完成！")
                        except Exception as e:
                            st.error(f"❌ 预测失败: {str(e)}")
        
        # 显示预测结果
        if st.session_state.get('prediction_results'):
            results = st.session_state.prediction_results
            
            # 预测摘要
            col1, col2, col3 = st.columns(3)
            with col1:
                st.metric("当前价格", f"¥{results['stats']['current_price']:.2f}")
            with col2:
                st.metric(f"第{predict_days}天预测", 
                         f"¥{results['future_prices'][-1]:.2f}",
                         f"{results['stats']['total_change']:+.2f}")
            with col3:
                st.metric("投资建议", results['suggestion'])
            
            # 预测图表
            prediction_chart = plot_prediction_results(df, results, f"{st.session_state.current_stock['name']}预测")
            st.plotly_chart(prediction_chart, use_container_width=True)
            
            # 下载数据
            csv = results['predictions'].to_csv(index=False)
            st.download_button(
                "📥 下载预测数据",
                csv,
                f"{current_stock_code}_{current_stock_name}_prediction.csv",
                "text/csv"
            )

else:
    st.info("💡 请输入股票代码或选择热门股票开始分析")

# 初始化session state
if 'model_trained' not in st.session_state:
    st.session_state.model_trained = False

# 页脚
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #666; padding: 1rem;'>
    <p>⚠️ <strong>风险提示</strong>: 预测结果仅供学习参考，投资有风险，决策需谨慎！</p>
    <p>🔧 支持全部A股代码 | 基于akshare免费数据源 | 机器学习预测</p>
</div>
""", unsafe_allow_html=True)